#include "NalParser.h"
//#include "H264NalParser.hpp"
//#include "H265NalParser.hpp"
#include <math.h>
#include "h264_stream.h"
#include "h265_stream.h"


int H264ReadNalUnit(h264_stream_t* h, unsigned char*buf,unsigned int nLen)
{
    return read_nal_unit(h, buf, nLen);
}

int H265ReadNalUnit(h265_stream_t* h, unsigned char*buf,unsigned int nLen)
{
    return h265_read_nal_unit(h, buf, nLen);
}


/*********************************************** 
*Description : 获取分辨率
               参考：https://latelee.blog.csdn.net/article/details/48844915
                     https://blog.csdn.net/subfate/article/details/48576445
*Input		 : code_name: 编码类型名(H264/H265)  buf: 帧数据缓冲区(不需要开始字符串) nLen: 帧数据长度
*Output		 : width：宽  height: 高
*Return      :
*Author      : jixiang.zhang
*data        : 2022-09-19
***********************************************/
void GetResolution(const char *code_name, unsigned char*buf, unsigned int nLen, unsigned int* width, unsigned int* height )
{
    int ret  = 0;
    if( strcmp(code_name, "H264") == 0)
    {
        h264_stream_t* m_hH264 = h264_new();
    
        ret = H264ReadNalUnit(m_hH264, buf, nLen);

        *width = (unsigned int)m_hH264->info->width;
        *height = (unsigned int)m_hH264->info->height;

        //printf("***file:%s function:%s line:%d H264 frame rate = %f\n", __FILE__, __FUNCTION__, __LINE__, m_hH264->info->max_framerate);

        if (m_hH264 != NULL)
        {
            h264_free(m_hH264);
            m_hH264 = NULL;
        }
    }
    else if(strcmp(code_name, "H265") == 0)
    {

        h265_stream_t* m_hH265 = h265_new();
    
        ret = H265ReadNalUnit(m_hH265, buf, nLen);

        *width = (unsigned int)m_hH265->info->width;
        *height = (unsigned int)m_hH265->info->height;


        if (m_hH265 != NULL)
        {
            h265_free(m_hH265);
            m_hH265 = NULL;
        }
    }
}


void GetSEIPayload (const char *code_name, unsigned char*buf, unsigned int nLen, std::vector<int8_t>& payload) {
  int ret  = 0;
  if( strcmp(code_name, "H264") == 0) {
    h264_stream_t* m_hH264 = h264_new();

    ret = H264ReadNalUnit(m_hH264, buf, nLen);
    //printf("payloadType = %d\n", m_hH264->sei->payloadType);
    //printf("payload = %s\n", m_hH264->sei->payload);
    //printf("payloadSize = %d\n", m_hH264->sei->payloadSize);
    for (int i = 0; i < m_hH264->sei->payloadSize; i++) {
      payload.push_back(m_hH264->sei->payload[i]);
    }

    //printf("***file:%s function:%s line:%d H264 frame rate = %f\n", __FILE__, __FUNCTION__, __LINE__, m_hH264->info->max_framerate);

    if (m_hH264 != NULL) {
        h264_free(m_hH264);
        m_hH264 = NULL;
    }
  }
  else if(strcmp(code_name, "H265") == 0) {
    printf("Error: file:%s function:%s line:%d\n", __FILE__, __FUNCTION__, __LINE__);
  }
}


 
/*********************************************** 
*Description : 获取帧率，获取的帧率是真实帧率的两倍，还没有研究透。
               参考：https://latelee.blog.csdn.net/article/details/48844915
                     https://blog.csdn.net/subfate/article/details/48576445
*Input		 : code_name: 编码类型名(H264/H265)  buf: 帧数据缓冲区(不需要开始字符串) nLen: 帧数据长度
*Output		 : framerate: 帧率
*Return      :
*Author      : jixiang.zhang
*data        : 2022-09-19
***********************************************/
void GetFrameRate(const char *code_name, unsigned char*buf, unsigned int nLen, float* framerate)
{
    int ret  = 0;
    if( strcmp(code_name, "H264") == 0)
    {
        h264_stream_t* m_hH264 = h264_new();
    
        ret = H264ReadNalUnit(m_hH264, buf, nLen);

        printf("file:%s function:%s line:%d H265 frame rate = %f\n", __FILE__, __FUNCTION__, __LINE__, m_hH264->info->max_framerate);

        *framerate = m_hH264->info->max_framerate;

        if (m_hH264 != NULL)
        {
            h264_free(m_hH264);
            m_hH264 = NULL;
        }
    }
    else if(strcmp(code_name, "H265") == 0)
    {
        h265_stream_t* m_hH265 = h265_new();
    
        ret = H265ReadNalUnit(m_hH265, buf, nLen);

        printf("file:%s function:%s line:%d H265 frame rate = %f\n", __FILE__, __FUNCTION__, __LINE__, m_hH265->info->max_framerate);

        *framerate = m_hH265->info->max_framerate;

        if (m_hH265 != NULL)
        {
            h265_free(m_hH265);
            m_hH265 = NULL;
        }
    }
}



unsigned int Ue(unsigned char *pBuff,unsigned int  nLen,unsigned int &nStartBit)
{
	unsigned int nZeroNum = 0;
	unsigned int dwRet = 0;
	while(nStartBit < nLen * 8)
	{
		if(pBuff[nStartBit/8]&(0x80 >> (nStartBit%8)))
			break;
		nZeroNum++;
		nStartBit++;
	}
	nStartBit++;

	for(unsigned int i=0;i<nZeroNum;i++)
	{
		dwRet <<= 1;
		if(pBuff[nStartBit/8]&(0x80 >> (nStartBit%8)))
		{
			dwRet += 1;
		}
		nStartBit++;
	}
		return (1 << nZeroNum)-1 + dwRet;
	
}

int Se(unsigned char *pBuff,unsigned int nLen,unsigned int &nStartBit)
{
	int UeVal = Ue(pBuff,nLen,nStartBit);
	double k = UeVal;
	int nValue = ceil(k/2);

	if(UeVal%2 ==0)
	    nValue =- nValue;
	
	return nValue;
}

unsigned int u(unsigned int BitCount,unsigned char* buf,unsigned int &nStartBit)
{
	unsigned int dwRet = 0;
	for(unsigned int i = 0;i < BitCount;i++)
	{
		dwRet <<= 1;
		if(buf[nStartBit / 8]&(0x80 >> (nStartBit % 8)))
		{
			dwRet += 1;
		}
		nStartBit++;
	}

	return dwRet;
}


// /* return h265 frame type: */
// /* 0x01 : I frame */
// /* 0x02 : P/B frame */
// /* 0x03 : SPS frame */
// /* 0x04 : PPS frame */
// /* 0x05 : SEI frame */
// /* 0x06 : VPS frame for H265 only */
// /* 0x07 : Other frame */
int h265_frame_type(unsigned char*buf,unsigned int nLen)
{

    int nal_unit_type = (buf[0]&0x7E)>>1;

    if(( 16 <= nal_unit_type) && (nal_unit_type <= 21))
    {
        return 0x01;  /* I frame */
    }
    else if(( 0 <= nal_unit_type) && (nal_unit_type <= 9))
    {
        return 0x02;  /* P/B frame */
    }
    else if(nal_unit_type == 32)
    {
        return 0x06;  /* VPS frame */
    }
    else if(nal_unit_type == 33)
    {
        return 0x03;  /* SPS frame */
    }
    else if(nal_unit_type == 34)
    {
        return 0x04;   /* PPS frame */
    }
    else if(nal_unit_type == 39)
    {
        return 0x05;   /* SEI frame */
    }
    else
    {
        return 0x07; /* Other frame */
    }
}

/* return h264 frame type: */
/* 0x01 : I frame */
/* 0x02 : P/B frame */
/* 0x03 : SPS frame */
/* 0x04 : PPS frame */
/* 0x05 : SEI frame */
/* 0x07 : Other frame */
int h264_frame_type(unsigned char*buf,unsigned int nLen)
{
    unsigned int StartBit = 0;
    int forbidden_zero_bit = u(1,buf,StartBit);
    int nal_ref_idc = u(2,buf,StartBit);
    int nal_unit_type = u(5,buf,StartBit);

    if(nal_unit_type == 5) 
    {
        return 0x01;  /* I frame */
    }
    else if(nal_unit_type == 1)  
    {
        return 0x02;  /* P/B frame */
    }
    else if(nal_unit_type == 7) 
    {
        return 0x03;  /* SPS frame */
    }
    else if(nal_unit_type == 8)  
    {
        return 0x04;  /* PPS frame */
    }
    else if(nal_unit_type == 6)  /* SEI  */
    {
        return 0x05;   /* SEI frame */
    }
    else
    {
        return 0x07; /* Other frame */
    }
}